package com.stress.testing.plugin.apache.rocketmq.listener;

import com.stress.testing.common.TraceContext;
import com.stress.testing.common.TraceContextHolder;
import com.stress.testing.common.conf.AgentConfig;
import com.stress.testing.common.conf.PluginConfig;
import org.apache.rocketmq.client.consumer.listener.MessageListener;
import org.apache.rocketmq.common.message.MessageExt;

import java.util.Map;

public abstract class MessageListenerWrapper implements MessageListener {
    private AgentConfig agentConfig;
    private PluginConfig pluginConfig;

    public MessageListenerWrapper(AgentConfig agentConfig, PluginConfig pluginConfig) {
        this.agentConfig = agentConfig;
        this.pluginConfig = pluginConfig;
    }

    protected void consumeMessageBefore(MessageExt messageExt) {
        TraceContextHolder.removeTraceContext();

        /*String topic = messageExt.getTopic();
        String topicPrefix = pluginConfig.get(PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX_KEY,
                DEFAULT_PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX);

        if (!topic.startsWith(topicPrefix) &&
                !topic.startsWith(MixAll.DLQ_GROUP_TOPIC_PREFIX + topicPrefix) &&
                !topic.startsWith(MixAll.RETRY_GROUP_TOPIC_PREFIX + topicPrefix)) {
            return;
        }*/

        TraceContext traceContext = new TraceContext(agentConfig.getTracePropertyPrefix());
        Map<String, String> properties = messageExt.getProperties();
        for (Map.Entry<String, String> entry : properties.entrySet()) {
            if (entry.getKey().startsWith(agentConfig.getTracePropertyPrefix())) {
                traceContext.setProperty(entry.getKey(), entry.getValue());
            }
        }

        // 可以简单的认为，如果msg的property中不包含trace_id,那么就不包含trace信息
        if (!traceContext.hasTraceID()) {
            return;
        }

        TraceContextHolder.setTraceContext(traceContext);
    }

    protected void consumeMessageAfter() {
        TraceContextHolder.removeTraceContext();
    }
}
